#!/usr/bin/env python3
import asyncio
import websockets
import json
import socket

def check_port():
    """检查8080端口是否开放"""
    try:
        sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        sock.settimeout(5)
        result = sock.connect_ex(('172.20.0.94', 8080))
        sock.close()
        if result == 0:
            print("✅ 端口8080开放")
            return True
        else:
            print("❌ 端口8080未开放")
            return False
    except Exception as e:
        print(f"❌ 检查端口失败: {e}")
        return False

async def test_websocket():
    """测试WebSocket连接"""
    uri = "ws://172.20.0.94:8080"
    try:
        print(f"正在测试WebSocket连接: {uri}")
        async with websockets.connect(uri, ping_interval=None, ping_timeout=None) as websocket:
            print("✅ WebSocket连接成功")
            
            # 发送简单的心跳消息
            heartbeat = {
                "type": "HEARTBEAT",
                "data": {
                    "timestamp": 1234567890
                }
            }
            
            await websocket.send(json.dumps(heartbeat))
            print("✅ 心跳消息发送成功")
            
            # 等待响应
            try:
                response = await asyncio.wait_for(websocket.recv(), timeout=5.0)
                print(f"✅ 收到响应: {response}")
            except asyncio.TimeoutError:
                print("⚠️ 等待响应超时")
            
    except Exception as e:
        print(f"❌ WebSocket连接失败: {e}")

async def main():
    print("=== WebSocket连接诊断 ===")
    
    # 检查端口
    if check_port():
        # 测试WebSocket
        await test_websocket()
    else:
        print("端口未开放，无法测试WebSocket连接")
    
    print("=== 诊断完成 ===")

if __name__ == "__main__":
    asyncio.run(main()) 